# 创建数据库url
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI


def chat(message: str):
    # 设置模型名称，模型密钥，模型地址
    model = ChatOpenAI(model="deepseek-chat",api_key= "你的密钥",base_url= "https://api.deepseek.com",temperature=0)
    url = "mysql+pymysql://root:123456@127.0.0.1:3306/test"
    # 创建数据库对象
    db = SQLDatabase.from_uri(url)
    # 数据库和模型结合
    sql_chain = create_sql_query_chain(model, db)
    # 模型根据问题生成sql语句,并去掉字符串开头的 "query = """ 部分和去掉字符串结尾的 """ 部分
    execute_query = sql_chain.invoke({"question": message})
    print(execute_query)

    # 判断开头并选择处理方式
    if execute_query.startswith("```sql"):
        # 如果开头是 "```sql"，执行第二种处理方式
        execute_query = execute_query.strip("```sql").rstrip("```")
    elif execute_query.startswith("SQLQuery:"):
        # 如果开头是 "SQLQuery:"，执行第三种处理方式
        execute_query = execute_query.strip("SQLQuery:").rstrip("```")
        execute_query = execute_query.replace("```sql", "")
    print(execute_query)

    # 创建数据库执行sql语句工具
    execute_sql_tool = QuerySQLDatabaseTool(db=db)
    # 执行sql语句, 返回结果
    resp = execute_sql_tool.invoke(execute_query)
    return resp